import os
from dotenv import load_dotenv
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser, JsonOutputParser

# --- 环境准备 (与前几课相同) ---
print("正在从 env.example 加载环境变量...")
load_dotenv(dotenv_path="env.example")

api_key = os.getenv("OPENAI_API_KEY")
model_name = os.getenv("OPENAI_MODEL", "qwen-plus-latest")

if not api_key:
    print("错误：请先在 env.example 文件中设置你的 OPENAI_API_KEY。")
else:
    print("环境变量加载成功！")
    # --- 核心代码开始 ---

    # --- 示例一：最简单的解析器 StrOutputParser ---
    print("\n--- 示例一：使用 StrOutputParser ---")
    
    # 1. 准备模型和提示
    chat = ChatOpenAI(model=model_name)
    prompt = ChatPromptTemplate.from_messages([
        ("system", "你是一位专业的作家。"),
        ("human", "请为关于{topic}的故事写一个吸引人的标题。")
    ])
    
    # 2. 定义一个简单的解析器
    # StrOutputParser 是最基础的解析器，它的作用就是简单地将模型返回的 AIMessage 对象中的 content 提取成字符串。
    output_parser = StrOutputParser()

    # 3. 创建链
    # 我们把三个组件串联起来：提示 -> 模型 -> 解析器
    chain = prompt | chat | output_parser
    
    # 4. 调用链
    response = chain.invoke({"topic": "一只会飞的企鹅"})
    
    print(f"主题: 一只会飞的企鹅")
    print(f"AI 生成的标题: {response}")
    print(f"结果的数据类型: {type(response)}") # 你会看到，结果是 str 类型，而不是 AIMessage


    # --- 示例二：获取结构化输出 JsonOutputParser ---
    print("\n--- 示例二：使用 JsonOutputParser ---")

    # 1. 准备模型 (复用之前的 chat 对象)
    
    # 2. 创建一个新的提示模板，这次我们有更复杂的要求
    # 我们在提示中明确要求 AI 输出 JSON 格式
    json_prompt = ChatPromptTemplate.from_template(
        """
        你是一个会讲笑话的机器人。
        根据用户的主题 `{topic}` 来生成一个笑话。
        请严格按照下面的 JSON 格式返回你的答案，不要有任何多余的解释：
        {{
            "setup": "这里是笑话的铺垫",
            "punchline": "这里是笑话的笑点"
        }}
        """
    )

    # 3. 定义 JSON 解析器
    json_parser = JsonOutputParser()

    # 4. 创建新的链
    json_chain = json_prompt | chat | json_parser
    
    # 5. 调用链
    joke_response = json_chain.invoke({"topic": "程序员"})
    
    print(f"主题: 程序员")
    print(f"AI 生成的笑话 (JSON 对象): {joke_response}")
    print(f"结果的数据类型: {type(joke_response)}") # 你会看到，结果是 dict 类型
    print(f"笑话铺垫: {joke_response['setup']}")
    print(f"笑话笑点: {joke_response['punchline']}")
    print("--------------------") 